﻿#pragma once

#include <learn/alg/alg_approximator.h>
#include <learn/alg/alg_geom.h>
#include <learn/alg/alg_iterator.h>
#include <learn/alg/alg_projection.h>

namespace xi
{

namespace alg
{

using namespace xi::math;
using namespace xi::geom;

template <std::size_t N> struct Alg_Curve_PIA : Alg_Iterator<std::shared_ptr<Geom_BSplineCurve<N>>>
{
    using Vec = VecT<N, double>; // P

    std::size_t degree{3};
    std::vector<Vec> data;
    std::vector<double> params;

  protected:
    double m_w;
    std::shared_ptr<Geom_BSplineCurve<N>> m_curve;

    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        if (params.empty())
            params = Alg_Parameter<N>::chord(data);

        auto knots = Alg_Knot<N>::average(params, degree);
        std::vector<Vec> poles(data.size());
        for (std::size_t i = 0; i < data.size(); i++)
            poles[i] = data[i];

        // 计算全正基配置矩阵的最小特征值
        MatrixXd A(data.size(), data.size());
        A.set_zero();
        for (std::size_t i = 0; i < A.rows(); i++)
        {
            auto span = find_span(knots.data(), knots.size(), degree, params[i]);
            auto basis = basis_funcs(knots.data(), degree, span, params[i]);
            for (std::size_t j = 0; j < degree + 1; j++)
                A(i, span - degree + j) += basis[j];
        }

        // 最佳权重 w = 2 / (1 + lambda(A)), 其中 lambda(A) \in (0,1) 是 A 的最小特征值
        for (std::size_t i = 0; i < A.rows(); i++)
            A(i, i) -= 1;
        m_w = 2 / (2 - A.eigenvalue());

        m_curve = std::make_shared<Geom_BSplineCurve<N>>(std::move(poles), std::move(knots), degree);
    }

    /**
     * @brief 迭代计算下一个曲线
     *
     * @return std::shared_ptr<Geom_BSplineCurve<N>>
     */
    std::shared_ptr<Geom_BSplineCurve<N>> next() override
    {
        std::vector<Vec> nPoles(data.size());
        for (std::size_t i = 0; i < nPoles.size(); i++)
            nPoles[i] = m_curve->pole(i) + (data[i] - m_curve->value(params[i])) * m_w;

        m_curve = std::make_shared<Geom_BSplineCurve<N>>(std::move(nPoles), m_curve->move_knots(), degree);
        return m_curve;
    }

    /**
     * @brief 迭代计算下一个曲线，保留前一条曲线的数据
     *
     * @return std::shared_ptr<Geom_BSplineCurve<N>>
     */
    std::shared_ptr<Geom_BSplineCurve<N>> next_preserve() override
    {
        std::vector<Vec> nPoles(data.size());
        for (std::size_t i = 0; i < nPoles.size(); i++)
            nPoles[i] = m_curve->pole(i) + (data[i] - m_curve->value(params[i])) * m_w;

        auto knots = m_curve->knots();
        m_curve = std::make_shared<Geom_BSplineCurve<N>>(std::move(nPoles), std::move(knots), degree);
        return m_curve;
    }

    /**
     * @brief 根据插值精度判断是否应该继续迭代
     *
     * @return true
     * @return false
     */
    bool should_continue() override
    {
        for (std::size_t i = 0; i < data.size(); i++)
        {
            auto t = Alg_Projection<N>::nearest(m_curve, data[i], params[i], Numeric::distance);
            if ((m_curve->value(t) - data[i]).modulus() > Numeric::distance)
                return true;
        }
        return false;
    }
};

template <std::size_t N> struct Alg_Curve_LSPIA : Alg_Iterator<std::shared_ptr<Geom_BSplineCurve<N>>>
{
    using Vec = VecT<N, double>; // P

    std::size_t degree{3};
    std::size_t poleNum;
    std::vector<Vec> data;
    std::vector<double> params;
    double error{1}; // 记录每次迭代的最大误差

  protected:
    std::shared_ptr<Geom_BSplineCurve<N>> m_curve;

    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        if (params.empty())
            params = Alg_Parameter<N>::chord(data);

        // 寻找最大曲率特征点
        std::vector<std::tuple<int, double>> curvatures;
        for (std::size_t i = 1; i < data.size() - 1; i++)
        {
            auto p1 = data[i - 1];
            auto p2 = data[i];
            auto p3 = data[i + 1];
            auto v1 = p2 - p1;
            auto v2 = p1 - p2 * 2 + p3;
            curvatures.emplace_back(i, cross(v1, v2).modulus() / std::pow(v1.modulus(), 3));
        }
        std::sort(curvatures.begin(), curvatures.end(),
                  [](auto &a, auto &b) { return std::get<1>(a) > std::get<1>(b); });
        std::sort(curvatures.begin(), curvatures.begin() + poleNum - 2,
                  [](auto &a, auto &b) { return std::get<0>(a) < std::get<0>(b); });

        std::vector<Vec> poles(poleNum);
        std::vector<double> params(poleNum);
        poles[0] = data.front();
        params[0] = params.front();
        for (std::size_t i = 1; i < poleNum - 1; i++)
        {
            auto j = std::get<0>(curvatures[i - 1]);
            poles[i] = data[j];
            params[i] = params[j];
        }
        poles[poleNum - 1] = data.back();
        params[poleNum - 1] = params.back();

        // 使用这些控制点对应的参数值来构造节点向量
        auto knots = Alg_Knot<N>::average(params, degree);
        m_curve = std::make_shared<Geom_BSplineCurve<N>>(std::move(poles), std::move(knots), degree);
    }

    /**
     * @brief 迭代计算下一个曲线
     *
     * @return std::shared_ptr<Geom_BSplineCurve<N>>
     */
    std::shared_ptr<Geom_BSplineCurve<N>> next() override
    {
        auto m = data.size() - 1;
        auto n = poleNum - 1;

        std::vector<double> basisSum(n + 1, 0); // 记录每个控制点对应的那组基函数和
        std::vector<Vec> deltaSum(n + 1);       // 记录每个控制点对应的那组偏差向量加权和

        // 对每个数据点计算增量
        auto knots = m_curve->knots();
        for (std::size_t i = 0; i < m + 1; i++)
        {
            auto delta = data[i] - m_curve->value(params[i]);

            // 计算此参数对应的非零基函数值
            auto span = find_span(knots.data(), knots.size(), degree, params[i]);
            auto basis = basis_funcs(knots.data(), degree, span, params[i]);

            // 累计参数
            for (std::size_t k = 0; k < degree + 1; k++)
            {
                basisSum[span - degree + k] += basis[k];
                deltaSum[span - degree + k] += delta * basis[k];
            }
        }

        // 平均基函数作为权重
        for (std::size_t i = 0; i < n + 1; i++)
            m_curve->set_pole(i, m_curve->pole(i) + deltaSum[i] / basisSum[i]);
        return m_curve;
    }

    /**
     * @brief 根据逼近精度判断是否应该继续迭代
     *
     * @return true
     * @return false
     */
    bool should_continue() override
    {
        double err = 0;
        std::size_t max = 0;
        for (std::size_t i = 0; i < data.size(); i++)
        {
            auto t = Alg_Projection<N>::nearest(m_curve, data[i], params[i], Numeric::distance);
            auto d = (m_curve->value(t) - data[i]).modulus();
            if (d > err)
            {
                err = d;
                max = i;
            }
        }
        // 限制不取右端点
        max = std::min(max, data.size() - 2);

        if (err > Numeric::distance && std::abs(error - err) > Numeric::distance)
        {
            error = err;

            // 在最大误差点所在区间插入节点，不能让控制点数超过数据点数
            if (m_curve->pole_num() < data.size())
            {
                // Q \in [u_{span}, u_{span+1}]
                auto span = find_span(m_curve->knots().data(), m_curve->knots().size(), m_curve->degree(), params[max]);
                auto upper =
                    std::lower_bound(params.begin() + max, params.end(), m_curve->knot(span + 1)) - params.begin();

                // 找到第一个大于等于 u_{span+1} 的参数位置
                std::vector<double> nknots;
                for (std::size_t i = max + 1; i < upper + 1; i++)
                {
                    double t = (params[i] + params[i - 1]) / 2;
                    m_curve->insert_knot(t);
                    poleNum++;

                    if (poleNum == data.size())
                        break;
                }
            }
            return true;
        }
        return false;
    }
};

template <std::size_t N> struct Alg_Curve_FairingPIA : Alg_Curve_LSPIA<N>
{
    using Vec = VecT<N, double>; // P

    std::vector<double> wStretch; // 每个控制点的拉伸权重
    std::vector<double> wBending; // 每个控制点的光顺权重

  protected:
    // 迭代格式 P' = A * P + B
    Eigen::SparseMatrix<double> m_A;
    Eigen::MatrixXd m_B;

    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        Alg_Curve_LSPIA<N>::init();

        // 计算控制点的导数和二阶导数
        auto B0 = Alg_CoeffMatrix::curve_least_square(this->params, this->m_curve->knots(), this->degree);
        auto [B1, B2] = Alg_CoeffMatrix::curve_stretch_bending(this->m_curve->knots(), this->degree);
        auto [A, B] = coeff_matrices(B0, B1, B2, this->wStretch, this->wBending, this->data);

        m_A = std::move(A);
        m_B = std::move(B);
    }

    /**
     * @brief 迭代计算下一个曲线
     *
     * @return std::shared_ptr<Geom_BSplineCurve<N>>
     */
    std::shared_ptr<Geom_BSplineCurve<N>> next() override
    {
        auto n = this->poleNum - 1;

        // 填充控制点矩阵
        Eigen::MatrixXd P(n + 1, N);
        for (std::size_t i = 0; i < n + 1; i++)
            for (std::size_t j = 0; j < N; j++)
                P(i, j) = this->m_curve->pole(i)[j];

        P = m_A * P + m_B;
        for (std::size_t i = 0; i < n + 1; i++)
        {
            Vec p;
            for (std::size_t j = 0; j < N; j++)
                p[j] = P(i, j);
            this->m_curve->set_pole(i, p);
        }

        return this->m_curve;
    }

    /**
     * @brief 根据插值精度判断是否应该继续迭代
     *
     * @return true
     * @return false
     */
    bool should_continue() override
    {
        return true;
    }

  public:
    /**
     * @brief 根据控制点的拟合、拉伸和光顺矩阵计算系数矩阵
     *
     * @param[in] B0
     * @param[in] B1
     * @param[in] B2
     * @param[in] wStretch
     * @param[in] wBending
     * @param[in] data
     * @return std::tuple<Eigen::SparseMatrix<double>, Eigen::MatrixXd>
     */
    static std::tuple<Eigen::SparseMatrix<double>, Eigen::MatrixXd> coeff_matrices(
        const Eigen::SparseMatrix<double> &B0, const Eigen::SparseMatrix<double> &B1,
        const Eigen::SparseMatrix<double> &B2, const std::vector<double> &wStretch, const std::vector<double> &wBending,
        const std::vector<Vec> &data)
    {
        auto n = wBending.size() - 1;
        std::vector<Eigen::Triplet<double>> tripletsO1, tripletsO2;
        for (int i = 0; i < n + 1; i++)
        {
            tripletsO1.push_back({i, i, wStretch[i]});
            tripletsO2.push_back({i, i, wBending[i]});
        }

        Eigen::SparseMatrix<double> O1(n + 1, n + 1), O2(n + 1, n + 1);
        O1.setFromTriplets(tripletsO1.begin(), tripletsO1.end());
        O2.setFromTriplets(tripletsO2.begin(), tripletsO2.end());

        // 计算系数矩阵
        Eigen::SparseMatrix<double> I(n + 1, n + 1);
        I.setIdentity();

        Eigen::SparseMatrix<double> A = (I - O1 - O2) * B0.transpose() * B0 + O1 * B1 + O2 * B2;

        // 构造权重 Lambda
        std::vector<Eigen::Triplet<double>> tripletsLam;

        // 计算行元素绝对值的和
        std::vector<double> rowSums(A.rows(), 0);
        for (int k = 0; k < A.outerSize(); ++k)
        {
            for (Eigen::SparseMatrix<double>::InnerIterator it(A, k); it; ++it)
                rowSums[it.row()] += std::abs(it.value());
        }
        // 这里必须用 int，不然类型不对
        for (int i = 0; i < rowSums.size(); i++)
            tripletsLam.push_back({i, i, 1 / rowSums[i]});

        Eigen::SparseMatrix<double> Lam(n + 1, n + 1);
        Lam.setFromTriplets(tripletsLam.begin(), tripletsLam.end());

        A = I - Lam * A;
        A.makeCompressed();

        // 填充点云矩阵
        Eigen::MatrixXd Q(data.size(), N);
        for (std::size_t k = 0; k < data.size(); k++)
            for (std::size_t j = 0; j < N; j++)
                Q(k, j) = data[k][j];

        // 最终的迭代矩阵
        return std::make_tuple(A, Lam * (I - O1 - O2) * B0.transpose() * Q);
    }
};

} // namespace alg

} // namespace xi
